Skip to content

[Runtime] Fix PagedKVCache for PopN and enhance tests#17045

Merged
tqchen merged 1 commit intoapache:mainfrom
MasterJH5574:tvm-dev/2024-05-29-kv-cache-removal-fix
May 30, 2024
Merged

[Runtime] Fix PagedKVCache for PopN and enhance tests#17045
tqchen merged 1 commit intoapache:mainfrom
MasterJH5574:tvm-dev/2024-05-29-kv-cache-removal-fix

Conversation

@MasterJH5574
Copy link
Copy Markdown
Contributor

This PR fixes a bug in the PagedKVCache which may happen when the sequence removal order is not consistent with the reverse order of sequence add/fork order. With this fix, the PagedKVCache now supports removing sequences in any order without breaking.

This PR also adds an empty function to PagedKVCache to check if the KV cache is empty. Right now this function is only used for test purpose, where we check if everything in the KV cache is freed after removing all sequences.

@MasterJH5574 MasterJH5574 requested a review from cyx-6 May 29, 2024 21:30
Copy link
Copy Markdown
Contributor

@cyx-6 cyx-6 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks perfect!

@MasterJH5574 MasterJH5574 force-pushed the tvm-dev/2024-05-29-kv-cache-removal-fix branch from d19daf4 to d052c2d Compare May 30, 2024 01:47
This PR fixes a bug in the PagedKVCache which may happen when the
sequence removal order is not consistent with the reverse order
of sequence add/fork order. With this fix, the PagedKVCache now
supports removing sequences in any order without breaking.

This PR also adds an `empty` function to PagedKVCache to check if
the KV cache is empty. Right now this function is only used for test
purpose, where we check if everything in the KV cache is freed after
removing all sequences.
@MasterJH5574 MasterJH5574 force-pushed the tvm-dev/2024-05-29-kv-cache-removal-fix branch from d052c2d to 46237a1 Compare May 30, 2024 13:29
@tqchen tqchen merged commit 1eac178 into apache:main May 30, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants